package org.apache.commons.httpclient.contrib.ssl;

import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;

import javax.net.ssl.SSLPeerUnverifiedException;
import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.security.cert.X509Certificate;

import org.apache.commons.httpclient.ConnectTimeoutException;
import org.apache.commons.httpclient.params.HttpConnectionParams;
import org.apache.commons.httpclient.protocol.SecureProtocolSocketFactory;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

public class StrictSSLProtocolSocketFactory implements SecureProtocolSocketFactory {

    private static final Log LOG            = LogFactory.getLog(StrictSSLProtocolSocketFactory.class);

    private boolean          verifyHostname = true;

    public StrictSSLProtocolSocketFactory(boolean verifyHostname){
        this.verifyHostname = verifyHostname;
    }

    public StrictSSLProtocolSocketFactory(){
    }

    public void setHostnameVerification(boolean verifyHostname) {
        this.verifyHostname = verifyHostname;
    }

    public boolean getHostnameVerification() {
        return this.verifyHostname;
    }

    public Socket createSocket(String host, int port, InetAddress clientHost, int clientPort) throws IOException,
                                                                                              UnknownHostException {
        SSLSocketFactory sf = (SSLSocketFactory) SSLSocketFactory.getDefault();
        SSLSocket sslSocket = (SSLSocket) sf.createSocket(host, port, clientHost, clientPort);

        verifyHostname(sslSocket);

        return sslSocket;
    }

    public Socket createSocket(String host, int port) throws IOException, UnknownHostException {
        SSLSocketFactory sf = (SSLSocketFactory) SSLSocketFactory.getDefault();
        SSLSocket sslSocket = (SSLSocket) sf.createSocket(host, port);
        verifyHostname(sslSocket);

        return sslSocket;
    }

    public Socket createSocket(Socket socket, String host, int port, boolean autoClose) throws IOException,
                                                                                        UnknownHostException {
        SSLSocketFactory sf = (SSLSocketFactory) SSLSocketFactory.getDefault();
        SSLSocket sslSocket = (SSLSocket) sf.createSocket(socket, host, port, autoClose);

        verifyHostname(sslSocket);

        return sslSocket;
    }

    private void verifyHostname(SSLSocket socket) throws SSLPeerUnverifiedException, UnknownHostException {
        if (!this.verifyHostname) {
            return;
        }
        SSLSession session = socket.getSession();
        String hostname = session.getPeerHost();
        InetAddress addr;
        try {
            addr = InetAddress.getByName(hostname);
        } catch (UnknownHostException uhe) {
            throw new UnknownHostException("Could not resolve SSL sessions server hostname: " + hostname);
        }

        X509Certificate[] certs = session.getPeerCertificateChain();
        if ((certs == null) || (certs.length == 0)) {
            throw new SSLPeerUnverifiedException("No server certificates found!");
        }

        String dn = certs[0].getSubjectDN().getName();

        if (LOG.isDebugEnabled()) {
            LOG.debug("Server certificate chain:");
            for (int i = 0; i < certs.length; i++) {
                LOG.debug("X509Certificate[" + i + "]=" + certs[i]);
            }
        }

        String cn = getCN(dn);
        if (hostname.equalsIgnoreCase(cn)) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Target hostname valid: " + cn);
            }
        } else {
            throw new SSLPeerUnverifiedException("HTTPS hostname invalid: expected '" + hostname + "', received '" + cn
                                                 + "'");
        }
    }

    private String getCN(String dn) {
        int i = 0;
        i = dn.indexOf("CN=");
        if (i == -1) {
            return null;
        }

        dn = dn.substring(i + 3);

        char[] dncs = dn.toCharArray();
        for (i = 0; i < dncs.length; i++) {
            if ((dncs[i] == ',') && (i > 0) && (dncs[(i - 1)] != '\\')) {
                break;
            }
        }
        return dn.substring(0, i);
    }

    @Override
    public Socket createSocket(String host, int port, InetAddress localAddress, int localPort,
                               HttpConnectionParams params) throws IOException, UnknownHostException,
                                                            ConnectTimeoutException {
        // TODO Auto-generated method stub
        return createSocket(host, localPort, localAddress, localPort);
    }
}
